import os
import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
import torch

class TxtDataLoader():
    def __init__(self, config):
        self.datasets_config = config["datasets"]
        self.label_type = self.datasets_config["label_type"]
        self.data_path = self.datasets_config["output_path"]
        if (self.data_path == None):
            self.data_path = "datasets/group"

    def get_train_loader(self, key):
        # 读取配置
        self.cfg_trainset = self.datasets_config["trainset"]
        bs = self.cfg_trainset["batch_size"]
        shuffle_ = self.cfg_trainset["shuffle"]

        self.datafile_path = self.data_path + "/" + str(key) + "/"
        files_name = os.listdir(self.datafile_path)
        assert len(files_name) == int(self.label_type)
        files_name.sort(key=lambda x:int(x.split(".")[0]))
        d_data = []
        d_label = []
        for file_name in files_name:
            file_path = self.datafile_path + file_name
            with open(file_path, 'r') as f:
                r_data = f.readlines()
            # 将r_data转为二维数组
            for i in range(len(r_data)):
                d_label.append(int(file_name.split(".")[0]))
                trace = r_data[i].strip().split()
                f_trace = []
                for t in range(len(trace)):
                    f_trace.append(float(trace[t]))
                d_data.append(f_trace)
        traces = torch.from_numpy(np.array(d_data)).float()
        labels = torch.from_numpy(np.array(d_label))
        dataset = TensorDataset(traces, labels)
        loader = DataLoader(dataset, batch_size=bs, shuffle=shuffle_)
        return loader
